#ifndef __KERNELRELOAD_KERNELRELOAD_KERNELLOADER_HPP__
#define __KERNELRELOAD_KERNELRELOAD_KERNELLOADER_HPP__

#include <ntimage.h>
#include "Fixrelocation.h"
namespace ntos
{
	template<typename element>
	class delete_helper_base
	{
	public:
		delete_helper_base()
			: pointer_(NULL)
		{}
		delete_helper_base(element* p)
			: pointer_(p)
		{}
		inline element* get() { return pointer_; }
		template<typename T>
		inline T* get() { return (T*)pointer_; }
		template<typename T>
		inline T* get(DWORD offset) { return (T*)((DWORD)pointer_ + offset); }
	protected:
		element* pointer_;
	};
	template<typename element, bool = false>
	class delete_helper : public delete_helper_base<element>
	{
	public:
		delete_helper()
		{}
		delete_helper(element* p)
			: delete_helper_base(p)
		{}
		~delete_helper() { NULL == pointer_ ? NULL : delete pointer_; }
	public:
		void operator = (element* p)
		{
			pointer_ = p;
		}
		void* operator + (DWORD offset)
		{
			return (void*)((DWORD)pointer_ + offset);
		}
		bool operator == (void* p)
		{
			return p == pointer_; 
		}
	};
	template<typename element>
	class delete_helper<element, true> : public delete_helper_base<element>
	{
	public:
		delete_helper()
		{}
		delete_helper(element* p)
			: delete_helper_base(p)
		{}
		~delete_helper() { NULL == pointer_ ? NULL : delete[] pointer_; }
	public:
		void operator = (element* p)
		{
			pointer_ = p;
		}
		void* operator + (DWORD offset)
		{
			return (void*)((DWORD)pointer_ + offset);
		}
		bool operator == (void* p)
		{
			return p == pointer_; 
		}
	};
	//
	// ObjectHolder
	//
	template<typename T>
	class ObjectHolder
	{
	public:
		ObjectHolder(T* p)
			:object_(p)
		{}
		~ObjectHolder()
		{ NULL != object_ ? delete object_ : NULL; }
		inline T* get() { return object_; }
		inline void attach(T* p) { object_ = p; }
		inline T* release() { T* p = object_; object_ = NULL; return p; }
		T* operator->()
		{
			return object_;
		}
		bool operator == (T* p)
		{
			return p == object_; 
		}
		bool operator != (T* p)
		{
			return p != object_; 
		}
	private:
		T* object_;
	};
	//
	// KernelReLoader
	//
	class KernelReLoader
	{
	public:
		KernelReLoader()
			: module_base_(NULL), module_size_(0)
		{
			RtlZeroMemory(module_name_, sizeof(module_name_));
		}
		~KernelReLoader()
		{
			if (NULL != module_base_)
				UnloadKernelFile();
		}
	public:
		template<typename Pool>
		BOOL LoadKernelFile(const char* file_name)
		{
			if (NULL == file_name)
				return FALSE;
			
			RtlCopyMemory(module_name_, file_name, strlen(file_name));

			HANDLE file_handle = OpenFile(file_name, GENERIC_READ, FILE_SHARE_READ);
			if (NULL == file_handle)
				return FALSE;

			DWORD file_size = GetModuleFileSize(file_handle);
			if (0 == file_size)
				return FALSE;

			delete_helper<BYTE, true> file_buffer(new(PagedPool) BYTE[file_size]);
			if (NULL == file_buffer.get<PVOID>())
				return FALSE;

			RtlZeroMemory(file_buffer.get<PVOID>(), file_size);

			if (!ReadFile(file_handle, file_buffer.get<PVOID>(), file_size))
				return FALSE;

			IMAGE_DOS_HEADER* dos_header = file_buffer.get<IMAGE_DOS_HEADER>();
			if (IMAGE_DOS_SIGNATURE != dos_header->e_magic)
				return FALSE;

			IMAGE_NT_HEADERS* nt_header = file_buffer.get<IMAGE_NT_HEADERS>(dos_header->e_lfanew);
			if (IMAGE_NT_SIGNATURE != nt_header->Signature)
				return FALSE;

			DWORD image_size = nt_header->OptionalHeader.SizeOfImage;
			DWORD file_alignment = nt_header->OptionalHeader.FileAlignment;
			DWORD section_alignment = nt_header->OptionalHeader.SectionAlignment;
			DWORD module_size_ = align_size(image_size, section_alignment);
			module_base_ = new((POOL_TYPE)Pool::pool_type) BYTE[module_size_];
			if (NULL == module_base_)
				return FALSE;
			
			RtlZeroMemory(module_base_, module_size_);

			IMAGE_SECTION_HEADER* section_header = file_buffer.get<IMAGE_SECTION_HEADER>(dos_header->e_lfanew + sizeof(IMAGE_NT_HEADERS));
			for (int i = 0; i < nt_header->FileHeader.NumberOfSections; ++i)
			{
				section_header[i].SizeOfRawData = align_size(section_header[i].SizeOfRawData, file_alignment);
				section_header[i].Misc.VirtualSize = align_size(section_header[i].Misc.VirtualSize, section_alignment);

				RtlCopyMemory(&module_base_[section_header[i].VirtualAddress], &file_buffer.get<BYTE>()[section_header[i].PointerToRawData], section_header[i].SizeOfRawData);
			}

			RtlCopyMemory(module_base_, file_buffer.get<BYTE>(), nt_header->OptionalHeader.SizeOfHeaders);

			return TRUE;
		}
		VOID UnloadKernelFile()
		{
			delete[] module_base_;
			module_base_ = NULL;
			module_size_ = 0;
			RtlZeroMemory(module_name_, sizeof(module_name_));
		}
	private:
		inline DWORD align_size(DWORD size, DWORD align) { return ((size + align - 1) / align * align); }

		HANDLE OpenFile(const char* file_name, ACCESS_MASK desired_access, ULONG share_access)
		{
			ANSI_STRING ansi_name;
			RtlInitAnsiString(&ansi_name, file_name);
			UNICODE_STRING unicode_name;
			if (!NT_SUCCESS(RtlAnsiStringToUnicodeString(&unicode_name, &ansi_name, TRUE)))
				return NULL;

			OBJECT_ATTRIBUTES object_att = { 0 };
			InitializeObjectAttributes(&object_att, &unicode_name, OBJ_CASE_INSENSITIVE, NULL, NULL);
			HANDLE file_handle = NULL;
			IO_STATUS_BLOCK io_status = { 0 };
			if (!NT_SUCCESS(ZwOpenFile(&file_handle, desired_access, &object_att, &io_status, share_access, FILE_SYNCHRONOUS_IO_NONALERT)))
				return NULL;

			return file_handle;
		}
		DWORD GetModuleFileSize(HANDLE file_handle)
		{
			IO_STATUS_BLOCK io_status = { 0 };
			FILE_STANDARD_INFORMATION file_info = { 0 };
			if (!NT_SUCCESS(ZwQueryInformationFile(file_handle, &io_status, &file_info, sizeof(FILE_STANDARD_INFORMATION), FileStandardInformation)))
				return 0;
			
			return file_info.EndOfFile.QuadPart;
		}
		BOOL ReadFile(HANDLE file_handle, PVOID buffer, DWORD read_size)
		{
			if (NULL == buffer)
				return FALSE;

			IO_STATUS_BLOCK io_status = { 0 };
			return NT_SUCCESS(ZwReadFile(file_handle, NULL, NULL, NULL, &io_status, buffer, read_size, NULL, NULL));
		}
		VOID CloseFile(HANDLE file_handle)
		{
			ZwClose(file_handle);
		}
	public:
		BOOL FixModuleImportTable(PVOID os_used_module_base)
		{
			if (NULL == os_used_module_base)
				return FALSE;

			DWORD iat_size = 0;
			IMAGE_THUNK_DATA* first_thunk = (IMAGE_THUNK_DATA*)RtlImageDirectoryEntryToData(module_base_, TRUE, IMAGE_DIRECTORY_ENTRY_IAT, &iat_size);

			DWORD orginel_iat_size = 0;
			IMAGE_THUNK_DATA* orginel_first_thunk = (IMAGE_THUNK_DATA*)RtlImageDirectoryEntryToData(os_used_module_base, TRUE, IMAGE_DIRECTORY_ENTRY_IAT, &orginel_iat_size);
			if (iat_size != orginel_iat_size)
			{
				KdPrint(("FixModuleImportTable >> iat_size != orginel_iat_size!!\r\n"));
				return FALSE;
			}
			RtlCopyMemory(first_thunk, orginel_first_thunk, iat_size);

			return TRUE;
		}
		BOOL FixModuleRelocTable(PVOID os_used_module_base)
		{
			if (NULL == os_used_module_base)
				return FALSE;

			return FixBaseRelocTable(module_base_, (DWORD)os_used_module_base);
		}
		inline BYTE* get_module_base() { return module_base_; }
		inline DWORD get_module_size() { return module_size_; }
	private:
		char module_name_[260];
		BYTE* module_base_;
		DWORD module_size_;
	};
}// namespace ntos
#endif//__KERNELRELOAD_KERNELRELOAD_KERNELLOADER_HPP__